from typing import Callable, TypeVar

import torch

T = TypeVar("T")
Compiling = bool | str | Callable


def resolve_compiling(compiling: Compiling) -> Callable[[T], T]:
    if compiling:
        if isinstance(compiling, bool):
            return torch.compile
        if isinstance(compiling, str):
            return torch.compile(backend=compiling)
        if isinstance(compiling, Callable):
            return compiling
        raise TypeError(f"Invalid compile option: {compiling}")
    return lambda arg: arg
